"""
Colossal AI 监控工具，用于监控大模型训练过程中的梯度和其他指标。
"""
import os

__version__ = "0.1.0"

# 设置环境变量
if "COLO_MONITOR_OUTPUT_DIR" not in os.environ:
    os.environ["COLO_MONITOR_OUTPUT_DIR"] = "./colo_monitor_output"

# 导入核心模块
from .core.monitor import TrainerMon, MonitorConfig

# 导入钩子模块
from .hooks import (
    GradientHook, ActivationHook, OptimizerHook, HookManager
)

# 导入指标计算模块
from .metrics.stats import (
    calculate_statistics, calculate_norm, calculate_max, calculate_min, calculate_mean
)

# 导入工具模块
from .utils.io import csv_to_tensorboard

__all__ = [
    # 版本信息
    "__version__",
    
    # 核心模块
    "TrainerMon",
    "MonitorConfig",
    
    # 钩子
    "GradientHook",
    "ActivationHook",
    "OptimizerHook",
    "HookManager",
    
    # 统计函数
    "calculate_statistics",
    "calculate_norm",
    "calculate_max",
    "calculate_min",
    "calculate_mean",
    
    # 工具函数
    "csv_to_tensorboard"
]
